import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

data = np.loadtxt(r'../data/test.txt')

data = StandardScaler().fit_transform(data)

x = data[:, 0]
y = data[:, 1]

plt.figure(figsize=[8, 8])
spr = 2
spc = 2
spn = 0

# plot data
spn += 1
plt.subplot(spr, spc, spn)
plt.scatter(x, y)

# elbow
spn += 1
plt.subplot(spr, spc, spn)
n_k = 10
ks = range(1, n_k + 1)
js = np.zeros(len(ks))
for i, k in enumerate(ks):
    km = KMeans(k)
    km.fit(data)
    j = km.inertia_
    js[i] = j
plt.plot(ks, js)
for i, k in enumerate(ks):
    plt.annotate(str(k), xy=[k, js[i]])

# clustering
spn += 1
plt.subplot(spr, spc, spn)
n_cluster = 4
km = KMeans(n_cluster)
km.fit(data)
h = km.labels_
cmap = plt.cm.get_cmap('rainbow', n_cluster)
plt.scatter(x, y, s=1, zorder=100, c=h, cmap=cmap)
centers = km.cluster_centers_
for i in range(n_cluster):
    c = cmap(i)
    plt.scatter(centers[i, 0], centers[i, 1], marker='x', s=100, zorder=0, color=c)
    plt.annotate('Cluster ' + str(i + 1), xy=[centers[i, 0], centers[i, 1]], color=c)

plt.show()
